from document_parser import extract_text
from collections import defaultdict
import os
import re


def create_index(directory):
    """
    为指定目录下的所有文件建立倒排索引。
    参数:
    directory: 要索引文件所在的目录路径。

    返回:
    index: 倒排索引字典，关键词映射到它们出现的文件和位置。
    supported_files: 目录中成功处理的文件列表。
    """
    # 使用 defaultdict 创建索引，列表用于存储关键词出现的位置
    index = defaultdict(list)
    # 构建文件路径列表
    supported_files = []
    # 获取目录下所有条目
    items = os.listdir(directory)
    # 遍历每个条目，检查是否为文件
    for item in items:
        file_path = os.path.join(directory, item)
        if os.path.isfile(file_path):
            supported_files.append(file_path)
            try:
                # 提取文件中的文本并转换为小写
                text = extract_text(file_path).lower()
                # 使用正则表达式找出所有的单词
                words = re.findall(r'\b\w+\b', text)
                # 遍历单词及其位置
                for pos, word in enumerate(words):
                    # 在倒排索引中记录单词出现的文件ID和位置
                    index[word].append((len(supported_files) - 1, pos))
            except Exception as e:
                print(f"Error processing file {file_path}: {e}")
        else:
            print(f"Skipping non-file item: {file_path}")
    return index, supported_files

def search(query, index, files):
    """
    根据查询词在索引中搜索，并返回匹配的文件列表。
    参数:
    query: 用户输入的查询词。
    index: 倒排索引字典。
    files: 文件列表。

    返回:
    results: 包含查询词的文件集合。
    """
    # 将查询词转换为小写
    query = query.lower()
    # 初始化结果集合
    results = set()
    # 如果查询词在索引中，获取匹配的文件ID集合
    if query in index:
        file_ids = set([file_id for file_id, _ in index[query]])
        # 根据文件ID获取文件名，添加到结果集
        for file_id in file_ids:
            results.add(files[file_id])
    return results
